/*
    Copyright 2007-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Implementation and explicit instantiations of the thick_ir_stage
/// class. \ingroup mcrx

// $Id$

#include "config.h"
#include "mcrx-stage.h"
#include "mcrx-stage-impl.h"
#include "create_grid.h"
#include "arepo_grid.h"
#include "arepo_emission.h"

using namespace std;

template <template <typename> class grid_type> 
bool mcrx::thick_ir_stage<grid_type>::check_stage_state () const
{
  bool state = false;
  try {
    this->output_file_->pHDU().readKey("MCRX" + stage_ID(), state);
  }
  catch (...) {}
  return state;
}

template <template <typename> class grid_type> 
void mcrx::thick_ir_stage<grid_type>::set_stage_state (bool state)
{
  this->output_file_->pHDU().addKey("MCRX" + stage_ID(), state, "");
}

template <template <typename> class grid_type> 
void mcrx::thick_ir_stage<grid_type>::basic_setup ()
{
  // read original and intensity wavelength vectors
  CCfits::ExtHDU& lambda_hdu = open_HDU (*this->output_file_, "LAMBDA");
  read(lambda_hdu.column("lambda"), elambda_);
  read(lambda_hdu.column("lambda_intensity"), ilambda_);
  int n_lambda_int;
  lambda_hdu.readKey("NLAMBDA_INTENSITY", n_lambda_int);
  ilambda_.resizeAndPreserve(n_lambda_int);

  // Read dust grains
  cout << "Setting up dust model" << endl;
  typename T_dust_model::T_scatterer_vector sv(read_dust_grains<typename T_dust_model::T_scatterer_vector>(this->p_, this->m_.units));
  this->model_.reset(new T_dust_model (sv));
  this->model_->set_wavelength(elambda_);

  // create shadow vector of grain_model* for temp calculation
  for(int i=0; i< this->model_->n_scatterers(); ++i) {
    grain_models_.push_back
      (dynamic_cast<grain_model<polychromatic_scatterer_policy, 
       mcrx_rng_policy>*>(&this->model_->get_scatterer(i)));
    grain_models_.back()->resample(ilambda_, elambda_);
  }

  // Now build grid
  cout << "Building grid" << endl;
  
  // The density generator is used to generate the densities of the
  // dust grains from the gas and metal densities in the simulation.  
  uniform_density_generator dg(this->p_);

  // we don't track intensities in this stage, so we can set n_lambda
  // to zero to save memory here.
  const int n_lambda = 0; 
 T_grid* dummy=0;
  this->g_ = create_grid(*this->input_file_, dg, n_lambda, dummy);

  // Create ir_grid emission object which calculates the dust emission
  // in the grid cells (temp calc is not done here)
  cout << "Setting up emission" << endl;

  // For the emission object we do however need to supply n_lambda.
  T_emission* dummy2=0;
  this->emi_ = create_grid(*this->input_file_, dg, elambda_.size(), dummy2);

  // set up cameras
  this->eme_.reset(new T_emergence (*this->output_file_));

  // set reference wavelength fixed for now
  vector<T_float> templambda(elambda_.begin(), elambda_.end());
  reflambda_= 
    lower_bound (templambda.begin(), templambda.end(),
		 this->p_.getValue("ir_reference_wavelength", T_float())) -
    templambda.begin() - 1;
  cout << "Reference wavelength is " << reflambda_ << endl;

  // read intensity data
  {
    cout << "Reading radiation intensities" << endl;
    CCfits::ExtHDU& intensity_hdu = open_HDU(*this->output_file_, "INTENSITY");
    read(intensity_hdu, intensity_);
    cout << "Allocating a " << intensity_.shape() 
       << " memory block for radiation intensities," 
       << product(intensity_.shape())*sizeof(T_float)*1.0/(1024*1024*1024)
       << " GB.\n";
    intensity_= pow (10., intensity_);

    // now add the dust intensity
    array_2 dust_intensity;
    CCfits::ExtHDU& dust_intensity_hdu = open_HDU(*this->output_file_, 
						  "DUSTINTENSITY");
    read(dust_intensity_hdu, dust_intensity);
    assert(all(dust_intensity.shape()==intensity_.shape()));
    intensity_ += pow(10., dust_intensity);
  }
}

template <template <typename> class grid_type> 
void mcrx::thick_ir_stage<grid_type>::setup_objects ()
{
  basic_setup();
}



/** Load IR stage data from a previously completed file. This is
     actually not supported. Because the temperature calculation have
     to be redone to convergence, there is very little lost by
     rerunning all rays. */
template <template <typename> class grid_type> 
void mcrx::thick_ir_stage<grid_type>::load_file ()
{
  basic_setup();
}

/** Load optically thick IR stage data from a dump file.  */ 
template <template <typename> class grid_type> 
void mcrx::thick_ir_stage<grid_type>::load_dump (binifstream& dump_file)
{
  basic_setup();

  if (!this->m_.terminator()) {
    // load current images
    this->eme_->load_dump(dump_file);
  }
}


/** Save images of IR emission.  */
template <template <typename> class grid_type> 
void mcrx::thick_ir_stage<grid_type>::save_file ()
{
  // Get normalization and save images
  const T_float normalization = this->emi_->total_emission_weight()/this->n_rays_;
  this->eme_->write_images(*this->output_file_, normalization, this->m_.units,
		     "-"+stage_ID(),
		     this->p_.defined("compress_images")&&
		     this->p_.getValue("compress_images", bool ()),
		     this->p_.defined("write_images_parallel")&&
		     this->p_.getValue("write_images_parallel", bool ()));

}


/** Dump IR stage info. We might have been interrupted during the temp
    calculation, or we might have been interrupted during shooting. In
    either case, we should save the cell SEDs that have been
    calculated so we don't have to do it again. */
template <template <typename> class grid_type> 
void mcrx::thick_ir_stage<grid_type>::save_dump (binofstream& dump_file) const
{
  this->eme_->write_dump(dump_file);
}


template <template <typename> class grid_type> 
template<typename T>
void mcrx::thick_ir_stage<grid_type>::calculate_dust_SED(const blitz::ETBase<T>& intensity)
{
  mcrx_terminator term(this->m_);
  const int n_threads = this->p_.defined("n_threads")?
    this->p_.getValue("n_threads", int ()): 1;

  const bool check_energy = this->p_.defined("check_energy_conservation") ?
    this->p_.getValue("check_energy_conservation", bool()) : false;

  this->emi_->calculate_SED(intensity, grain_models_, term, n_threads, 0,
		      false);
}


template <template <typename> class grid_type> 
bool mcrx::thick_ir_stage<grid_type>::shoot ()
{
  mcrx_terminator term(this->m_);

  // save forced scattering values and set them to no forced
  // scattering for this part, because it doesn't make sense.
  const int n_forced = this->p_.getValue("n_forced", int());
  const T_float i_forced = this->p_.getValue("i_forced", T_float());
  this->p_.setValue("n_forced",0);
  this->p_.setValue("i_forced",1);

  {
    // use weak references to avoid locking issues
    array_2 intensity_ref;
    intensity_ref.weakReference(intensity_);

    calculate_dust_SED(intensity_ref);
  }

  // at this point, we no longer need the intensity array, so we can free it
  intensity_.free();

  // we don't need to track intensities now, this is purely for
  // propagating to cameras, so add_intensity is false.
  this->m_.mon_.set_images(*this->eme_, reflambda_);
  bool t = this->shoot_scatter (false);
  this->m_.mon_.clear_images();

  this->p_.setValue("n_forced", n_forced);
  this->p_.setValue("i_forced", i_forced);

  return t;
}

template <template <typename> class grid_type> 
void mcrx::thick_ir_stage<grid_type>::operator() ()
{
  // this is here so the function is instantiated
  this->run_stage();
}


// explicit instantiations
template class mcrx::thick_ir_stage<mcrx::adaptive_grid>;

#ifdef WITH_AREPO
template class mcrx::thick_ir_stage<mcrx::arepo_grid>;
#endif
